from trinity.sceneRenderJobBase import SceneRenderJobBase
import trinity
import blue

def CreateSceneRenderJobSpace(name = None, stageKey = None):
    newRJ = SceneRenderJobSpace()
    if (name is not None):
        newRJ.ManualInit(name)
    else:
        newRJ.ManualInit()
    newRJ.SetMultiViewStage(stageKey)
    return newRJ



class SceneRenderJobSpace(SceneRenderJobBase):
    renderStepOrder = ['UPDATE_SCENE1',
     'UPDATE_SCENE',
     'UPDATE_UI',
     'SET_HDR_RT',
     'SET_INTZ_DEPTH',
     'SET_VAR_DEPTH',
     'SET_VIEWPORT',
     'SET_PROJECTION',
     'SET_VIEW',
     'CLEAR',
     'RENDER_SCENE',
     'SET_SCENE1_STATES',
     'RENDER_SCENE1',
     'RJ_POSTPROCESSING',
     'SET_FINAL_RT',
     'FINAL_BLIT',
     'POST_RENDER_CALLBACK',
     'UPDATE_TOOLS',
     'RENDER_INFO',
     'RENDER_VISUAL',
     'RENDER_TOOLS',
     'RENDER_UI']
    multiViewStages = [('SETUP',
      True,
      ['SET_VIEWPORT',
       'CLEAR',
       'UPDATE_SCENE1',
       'UPDATE_SCENE',
       'SET_INTZ_DEPTH',
       'SET_VAR_DEPTH']),
     ('SPACE_MAIN',
      False,
      ['SET_VIEWPORT',
       'SET_PROJECTION',
       'SET_VIEW',
       'RENDER_SCENE',
       'SET_SCENE1_STATES',
       'RENDER_SCENE1']),
     ('FINALIZE',
      True,
      ['RJ_POSTPROCESSING',
       'SET_FINAL_RT',
       'FINAL_BLIT']),
     ('TOOLS',
      False,
      ['SET_VIEWPORT',
       'SET_VIEW',
       'SET_PROJECTION',
       'UPDATE_TOOLS',
       'RENDER_INFO',
       'RENDER_VISUAL',
       'RENDER_TOOLS'])]
    visualizations = []
    renderTargetList = []

    def _ManualInit(self, name = 'SceneRenderJobSpace'):
        self.scene1 = None
        self.scene2 = None
        self.clientToolsScene = None
        self.activeSceneKey = None
        self.camera = None
        self.customBackBufferTexture = None
        self.intZTexture = None
        self.multiSampledRT = None
        self.depthMap = None
        self.depthMapMSSurface = None
        self.shadowMap = None
        self.postProcesses = []
        self.ui = None
        self.viewport = None
        self.hdrEnabled = False
        self.useBloom = False
        self.antiAliasing = False
        self.useShadows = False
        self.postProcessRJ = None



    def SetScene1(self, scene):
        if (scene is None):
            self.scene1 = None
        else:
            self.scene1 = blue.BluePythonWeakRef(scene)
        self.AddStep('UPDATE_SCENE1', trinity.TriStepUpdate(scene))
        self.AddStep('RENDER_SCENE1', trinity.TriStepRenderScene(scene))



    def GetScene1(self):
        if (self.scene1 is None):
            return 
        else:
            return self.scene1.object



    def SetClientToolsScene(self, scene):
        if (scene is None):
            self.clientToolsScene = None
        else:
            self.clientToolsScene = blue.BluePythonWeakRef(scene)
        self.AddStep('UPDATE_TOOLS', trinity.TriStepUpdate(scene))
        self.AddStep('RENDER_TOOLS', trinity.TriStepRenderScene(scene))



    def GetClientToolsScene(self):
        if (self.clientToolsScene is None):
            return 
        else:
            return self.clientToolsScene.object



    def SetUI(self, ui):
        if (ui is None):
            self.RemoveStep('UPDATE_UI')
            self.RemoveStep('RENDER_UI')
        else:
            self.AddStep('UPDATE_UI', trinity.TriStepUpdate(ui))
            self.AddStep('RENDER_UI', trinity.TriStepRenderUI(ui))



    def SetActiveCamera(self, camera):
        if (camera is None):
            self.RemoveStep('SET_VIEW')
            self.RemoveStep('SET_PROJECTION')
        else:
            self.AddStep('SET_VIEW', trinity.TriStepSetView(None, camera))
            self.AddStep('SET_PROJECTION', trinity.TriStepSetProjection(camera.triProjection))



    def SetActiveScenes(self, scene1, scene2, key):
        self.activeSceneKey = key
        self.SetScene1(scene1)
        self.SetScene(scene2)
        if (self.GetStep('RJ_POSTPROCESSING') is not None):
            self.AddStep('RJ_POSTPROCESSING', trinity.TriStepRunJob(self.CreatePostProcessRenderJob()))



    def _SetDepthMap(self):
        if (self.GetScene() is None):
            return 
        if hasattr(self.GetScene(), 'depthTexture'):
            self.GetScene().depthTexture = self.depthMap
            self.GetScene().multiSampledDepthSurface = self.depthMapMSSurface



    def _SetShadowMap(self):
        if (self.GetScene() is None):
            return 
        if self.useShadows:
            self.GetScene().shadowMap = self.shadowMap
        else:
            self.GetScene().shadowMap = None



    def EnablePostProcessing(self, enabled):
        if enabled:
            self.AddStep('RJ_POSTPROCESSING', trinity.TriStepRunJob(self.CreatePostProcessRenderJob()))
        else:
            self.RemoveStep('RJ_POSTPROCESSING')



    def EnablePostRenderCallbacks(self, enabled):
        if enabled:
            self.AddStep('POST_RENDER_CALLBACK', trinity.TriStepPostRenderCB())
        else:
            self.RemoveStep('POST_RENDER_CALLBACK')



    def AddPostProcess(self, id, path, sceneKey = None):
        entries = [ each for each in self.postProcesses if (each[0] == id) ]
        if (len(entries) > 0):
            return 
        self.postProcesses.append((id,
         path,
         sceneKey,
         {}))
        if (self.GetStep('RJ_POSTPROCESSING') is not None):
            self.AddStep('RJ_POSTPROCESSING', trinity.TriStepRunJob(self.CreatePostProcessRenderJob()))



    def RemovePostProcess(self, id):
        entries = [ each for each in self.postProcesses if (each[0] == id) ]
        for each in entries:
            self.postProcesses.remove(each)

        if (self.GetStep('RJ_POSTPROCESSING') is not None):
            self.AddStep('RJ_POSTPROCESSING', trinity.TriStepRunJob(self.CreatePostProcessRenderJob()))



    def SetPostProcessVariable(self, id, variable, value):
        ppSteps = getattr(self.postProcessRJ, 'steps', [])
        entries = [ each for each in ppSteps if (each.name == ('PostProcess ' + str(id))) ]
        for each in entries:
            for stage in each.PostProcess.stages:
                for param in stage.parameters:
                    if (param.name == variable):
                        param.value = value



        for each in self.postProcesses:
            if (each[0] == id):
                each[3][variable] = value
                break




    def SetupPostProcess(self, rj, id, pp, variables):
        if (self.customBackBufferTexture is not None):
            step = rj.RenderPostProcess(pp, self.customBackBufferTexture.GetSurfaceLevel(0), self.customBackBufferTexture)
            step.name = ('PostProcess ' + str(id))
        else:
            step = rj.RenderPostProcess(pp)
            step.name = ('PostProcess ' + str(id))
        for each in variables:
            self.SetPostProcessVariable(id, each, variables[each])




    def CreatePostProcessRenderJob(self):
        self.postProcessRJ = rj = trinity.CreateRenderJob('PostProcessingJob')
        bloomPP = trinity.device.postProcess
        if (bloomPP is None):
            return 
        else:
            if self.multiSampledRT:
                rj.CopySurfaceToRT(self.multiSampledRT, self.customBackBufferTexture.GetSurfaceLevel(0))
            self.SetupPostProcess(rj, 'Bloom', bloomPP, {})
            for (id, path, sceneKey, variables,) in self.postProcesses:
                if ((sceneKey == self.activeSceneKey) or (sceneKey is None)):
                    self.SetupPostProcess(rj, id, trinity.Load(path), variables)

            if self.multiSampledRT:
                rj.CopySurfaceToRT(self.customBackBufferTexture.GetSurfaceLevel(0), self.multiSampledRT)
            return rj



    def _SetScene(self, scene):
        self.currentMultiViewStageKey
        self.SetStepAttr('UPDATE_SCENE', 'object', scene)
        self.SetStepAttr('RENDER_SCENE', 'scene', scene)
        self._SetDepthMap()
        self._SetShadowMap()



    def _CreateBasicRenderSteps(self):
        self.AddStep('UPDATE_SCENE1', trinity.TriStepUpdate(self.GetScene1()))
        self.AddStep('RENDER_SCENE1', trinity.TriStepRenderScene(self.GetScene1()))
        self.AddStep('UPDATE_SCENE', trinity.TriStepUpdate(self.GetScene()))
        self.AddStep('RENDER_SCENE', trinity.TriStepRenderScene(self.GetScene()))
        self.AddStep('SET_SCENE1_STATES', trinity.TriStepSetStdRndStates(trinity.RM_ALPHA_ADDITIVE))
        self.AddStep('CLEAR', trinity.TriStepClear((0.0, 0.0, 0.0, 0.0), 1.0))
        if (self.clientToolsScene is not None):
            self.SetClientToolsScene(self.clientToolsScene.object)



    def DoReleaseResources(self, level):
        self.customBackBufferTexture = None
        self.multiSampledRT = None
        self.depthMap = None
        self.depthMapMSSurface = None
        self.viewport = None
        self.hdrEnabled = False
        self.useBloom = False
        self.antiAliasing = False
        self.useShadows = False
        self.shadowMap = None
        self.intZTexture = None
        self.SetRenderTargets(None, None, None, None, None)
        self.renderTargetList = None
        self.multiSampledRT = None
        self.customBackBufferTexture = None
        self.intZTexture = None
        self.depthMap = None
        self.depthMapMSSurface = None
        self._SetDepthMap()
        self._SetShadowMap()



    def NotifyResourceCreationFailed(self):
        try:
            eve.Message('CustomError', {'error': mls.UI_SHARED_VIDEOMEMORYERROR})

        finally:
            pass




    def _DoPrepareResources(self):
        self.viewport = trinity.device.viewport
        width = self.viewport.width
        height = self.viewport.height
        self.hdrEnabled = trinity.device.hdrEnable
        self.useBloom = (trinity.device.postProcess is not None)
        if sm.IsServiceRunning('device'):
            settings = sm.GetService('device').GetSettings()
            defaultShaderQuality = sm.GetService('device').GetDefaultShaderQuality()
            isShadowingSupported = sm.GetService('device').IsShadowingSupported()
            bbFormat = settings.BackBufferFormat
            dsFormat = settings.AutoDepthStencilFormat
            msType = settings.MultiSampleType
            msQuality = settings.MultiSampleQuality
        else:
            settings = trinity.device.GetPresentParameters()
            defaultShaderQuality = 3
            isShadowingSupported = True
            bbFormat = settings['BackBufferFormat']
            dsFormat = settings['AutoDepthStencilFormat']
            msType = settings['MultiSampleType']
            msQuality = settings['MultiSampleQuality']
        shaderQuality = prefs.GetValue('shaderQuality', defaultShaderQuality)
        if trinity.GetShaderModel().startswith('SM_3'):
            self.useShadows = prefs.GetValue('shadowsEnabled', isShadowingSupported)
        self.antiAliasing = (msType != trinity.TRIMULTISAMPLE_NONE)
        useDepth = trinity.GetShaderModel().endswith('DEPTH')
        self.depthMap = None
        if (useDepth and (not self.antiAliasing)):
            self.intZTexture = trinity.device.CreateTexture(width, height, 1, trinity.TRIUSAGE_DEPTHSTENCIL, trinity.TRIFMT_INTZ, trinity.TRIPOOL_DEFAULT)
        elif useDepth:
            self.depthMap = trinity.device.CreateTexture(width, height, 1, trinity.TRIUSAGE_RENDERTARGET, trinity.TRIFMT_R32F, trinity.TRIPOOL_DEFAULT)
            if self.antiAliasing:
                self.depthMapMSSurface = trinity.device.CreateRenderTarget(width, height, trinity.TRIFMT_R32F, msType, msQuality, False)
        if self.hdrEnabled:
            self.customBackBufferTexture = trinity.device.CreateTexture(width, height, 1, trinity.TRIUSAGE_RENDERTARGET, trinity.TRIFMT_A16B16G16R16F, trinity.TRIPOOL_DEFAULT)
            if self.antiAliasing:
                self.multiSampledRT = trinity.device.CreateRenderTarget(width, height, trinity.TRIFMT_A16B16G16R16F, msType, msQuality, False)
        if self.useShadows:
            self.shadowMap = trinity.TriShadowMap()
        self.renderTargetList = (blue.BluePythonWeakRef(self.multiSampledRT),
         blue.BluePythonWeakRef(self.customBackBufferTexture),
         blue.BluePythonWeakRef(self.intZTexture),
         blue.BluePythonWeakRef(self.depthMap),
         blue.BluePythonWeakRef(self.depthMapMSSurface))
        thingToSet = (x.object for x in self.renderTargetList)
        self.SetRenderTargets(*thingToSet)



    def DoPrepareResources(self):
        if ((not self.enabled) or (not self.canCreateRenderTargets)):
            return 
        try:
            self._DoPrepareResources()
        except trinity.D3DERR_OUTOFVIDEOMEMORY:
            log.LogException()
            self.DoReleaseResources(1)
            self.CreateRecoverySteps()
            uthread.new(self.NotifyResourceCreationFailed)



    def SetMultiViewStage(self, stageKey):
        self.currentMultiViewStageKey = stageKey



    def SetRenderTargets(self, multiSampledRT, customBackBufferTexture, intZTexture, depthMap, depthMapMSSurface):
        if (self.customBackBufferTexture is not None):
            if (self.multiSampledRT is not None):
                self.AddStep('SET_HDR_RT', trinity.TriStepSetRenderTarget(self.multiSampledRT))
                self.AddStep('FINAL_BLIT', trinity.TriStepCopySurfaceToRT(self.multiSampledRT, trinity.device.GetBackBuffer()))
            else:
                self.AddStep('SET_HDR_RT', trinity.TriStepSetRenderTarget(self.customBackBufferTexture.GetSurfaceLevel(0)))
                self.AddStep('FINAL_BLIT', trinity.TriStepCopySurfaceToRT(self.customBackBufferTexture.GetSurfaceLevel(0), trinity.device.GetBackBuffer()))
            self.AddStep('SET_FINAL_RT', trinity.TriStepSetRenderTarget(trinity.device.GetBackBuffer()))
        else:
            self.RemoveStep('SET_HDR_RT')
            self.RemoveStep('FINAL_BLIT')
            self.RemoveStep('SET_FINAL_RT')
        if (self.intZTexture is not None):
            self.AddStep('SET_INTZ_DEPTH', trinity.TriStepSetDepthStencil(self.intZTexture.GetSurfaceLevel(0)))
            self.AddStep('SET_VAR_DEPTH', trinity.TriStepSetVariableStore('DepthMap', self.intZTexture))
        else:
            self.RemoveStep('SET_INTZ_DEPTH')
            self.RemoveStep('SET_VAR_DEPTH')
        if depthMap:
            self._SetDepthMap()
        if (self.GetStep('RJ_POSTPROCESSING') is not None):
            self.AddStep('RJ_POSTPROCESSING', trinity.TriStepRunJob(self.CreatePostProcessRenderJob()))



    def GetRenderTargets(self):
        return self.renderTargetList




